import os
import shutil
import process_file as pf

path_list_split_group = '../dataset/split_group.json'
source_dir = "../result/TF_FUT/"
target_dir = ""

save_dir_root = "../result/json_train_valid_test/"
train_dir = f"{save_dir_root}train/"
valid_dir = f"{save_dir_root}valid/"
test_dir = f"{save_dir_root}test/"

len_train = 0
len_valid = 0
len_test = 0

list_split_group = pf.read_json(path_list_split_group)
os.makedirs(train_dir, exist_ok=True)
os.makedirs(valid_dir, exist_ok=True)
os.makedirs(test_dir, exist_ok=True)



len_list_split_group = len(list_split_group)
for index, group in enumerate(list_split_group):
    if index / len_list_split_group <= 0.8:
        target_dir = train_dir
    elif index / len_list_split_group <= 0.9:
        target_dir = valid_dir
    else:
        target_dir = test_dir
    for name in group:
        json_name = name + ".json"
        json_path = os.path.join(source_dir, json_name)
        if not os.path.isfile(json_path):
            continue
        list_dict_info = pf.read_json(json_path)

        if target_dir == train_dir:
            len_train += len(list_dict_info)
        elif target_dir == valid_dir:
            len_valid += len(list_dict_info)
        else:
            len_test += len(list_dict_info)

        if os.path.isfile(json_path):
            # 复制到目标目录
            shutil.copy(json_path, target_dir)
            print(f"{name} success!!!")
print(f"train data len: {len_train}")
print(f"valid data len: {len_valid}")
print(f"test data len: {len_test}")